package com.github.distributionmessage.integration.kafka;

import com.github.distributionmessage.constant.CommonConstant;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang.StringUtils;
import org.apache.kafka.clients.consumer.Consumer;
import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.springframework.kafka.support.Acknowledgment;
import org.springframework.kafka.support.converter.MessagingMessageConverter;
import org.springframework.messaging.Message;
import org.springframework.messaging.support.MessageBuilder;

import java.lang.reflect.Type;
import java.util.HashMap;
import java.util.Map;

@Slf4j
public class CustomKafkaMessageConverter extends MessagingMessageConverter {

    private String senderId;

    private String receiverId;

    private String serviceUrl;

    private String ieType;

    private String startNode;

    private String endNode;

    public CustomKafkaMessageConverter(String startNode, String endNode, String senderId, String receiverId, String serviceUrl, String ieType) {
        this.startNode = startNode;
        this.endNode = endNode;
        this.senderId = senderId;
        this.receiverId = receiverId;
        this.serviceUrl = serviceUrl;
        this.ieType = ieType;
    }

    @Override
    public Message<?> toMessage(ConsumerRecord<?, ?> record, Acknowledgment acknowledgment, Consumer<?, ?> consumer, Type type) {
        Message<?> originalMessage =  super.toMessage(record, acknowledgment, consumer, type);
        if (null == originalMessage) {
            return null;
        }

        Map<String, Object> originHeaders = originalMessage.getHeaders();
        Map<String, Object> headers = new HashMap<>(originHeaders);
        addCustomsHeaders(headers);

        return MessageBuilder.withPayload(originalMessage.getPayload())
                .copyHeaders(headers)
                .build();
    }

    public void addCustomsHeaders(Map<String, Object> source) {
        if (null != source) {
            if (StringUtils.isNotBlank(this.startNode)) {
                source.put(CommonConstant.START_NODE, this.startNode);
            }
            if (StringUtils.isNotBlank(this.endNode)) {
                source.put(CommonConstant.END_NODE, this.endNode);
            }
            if (StringUtils.isNotBlank(this.senderId)) {
                source.put(CommonConstant.SENDER_ID, this.senderId);
            }
            if (StringUtils.isNotBlank(this.receiverId)) {
                source.put(CommonConstant.RECEIVE_ID, this.receiverId);
            }
            if (StringUtils.isNotBlank(this.serviceUrl)) {
                source.put(CommonConstant.SIGN_AND_WRAP_SERVICE_URL, this.serviceUrl);
            }
            if (StringUtils.isNotBlank(this.ieType)) {
                source.put(CommonConstant.SIGN_AND_WRAP_IE_TYPE, this.ieType);
            }
        }
    }

    public static CustomKafkaMessageConverter createWrapMessageConverter(String senderId, String receiverId) {
        return new CustomKafkaMessageConverter(null, null, senderId, receiverId, null, null);
    }
    public static CustomKafkaMessageConverter createSvWrapMessageConverter(String startNode, String endNode) {
        return new CustomKafkaMessageConverter(startNode, endNode, null, null, null, null);
    }

    public static CustomKafkaMessageConverter createSignAndWrapMessageConverter(String serviceUrl, String ieType) {
        return new CustomKafkaMessageConverter(null, null, null, null, serviceUrl, ieType);
    }
}
